################################################################################

class rap_jconfig:
  """ Stores a single configuration in JJ-coupling. """
  def __init__(self,conforne,ne=None,ssocc=None):
    if isinstance(conforne,int):
      self.v_conf=self._baseconfig(conforne)
    else:
      self.v_conf=self._format_config(conforne,ne=ne,ssocc=ssocc)

# ------------------------------------------------------------------------------

  def __repr__(self):
    return self.label()

# ------------------------------------------------------------------------------

  def __eq__(self,conf):
    tconf=conf
    if not isinstance(tconf,rap_jconfig): tconf=rap_jconfig(tconf)
    if repr(self) == repr(tconf): return True
    return False

# ------------------------------------------------------------------------------

  def __ne__(self,conf):
    tconf=conf
    if not isinstance(tconf,rap_jconfig): tconf=rap_jconfig(tconf)
    if repr(self) != repr(tconf): return True
    return False

# ------------------------------------------------------------------------------

  def __cmp__(self,conf):
    return rlf_jconfsorter(self,conf)

# ------------------------------------------------------------------------------

  def __hash__(self):
    return hash(self.label())

# ------------------------------------------------------------------------------

  def __len__(self):
    """ Return length of configuration label. """
    return len(self.label())

# ------------------------------------------------------------------------------

  def __add__(self,other):
    return self.label()+other

# ------------------------------------------------------------------------------

  def __radd__(self,other):
    return other+self.label()

# ------------------------------------------------------------------------------

  def __getitem__(self,idx):
    return self.occupancy(idx)

# ------------------------------------------------------------------------------

  def __iter__(self):
    return self.subshells().__iter__()

# ------------------------------------------------------------------------------

  def _baseconfig(self,ne):
    """ Fill sub-shells in order until number of electrons is reached. """
    netot=ne
    out={}
    i=1
    while netot > 0:
      w=rlf_jssoccupancy(i)
      out[i]=min(w,netot)
      netot-=w
      i+=1
    return out

# ------------------------------------------------------------------------------

  def _format_config(self,inconf,ne=None,ssocc=None):
    """ Attempt to understand input configuration. """

    # configuration must be a string
    if isinstance(inconf,rap_jconfig):
      return inconf
    if not isinstance(inconf,str): 
      msg='configuration must be a string; conf: %s'%inconf
      raise ErrorRAP(msg)

    conf=inconf.lower()
    out={}

    # standard format
    tmp=re.findall(r'(\d+[a-z][+-])(\d+)',conf)
    if len(tmp) == 0:
      msg='cannot understand configuration string'
      raise ErrorRAP(msg)
    for (ss,tw) in tmp:
      w=int(tw)
      wmax=rlf_jssoccupancy(ss)
      if w > wmax:
        msg='occupancy exceeded for %s'%ss
        raise ErrorRAP(msg)
      iss=rlf_jssindex(ss)
      out[iss]=w

    # use subshell occupancies to figure out unspecified occupancies
    if ssocc is not None:
      for ss,ssw in ssocc.iteritems():
        (n,l)=rlf_ssdecode(ss)
        jss2=rlf_jssencode(n,l,+1,index=True)
        if l == 0:
          if jss2 not in out: out[jss2]=ssocc[ss]
          if out[jss2] != ssw:
            msg='configuration string incompatible with given sub-shell '+\
                'occupancy (%s)'%ss
            raise ErrorRAP(msg)
          continue
        jss1=rlf_jssencode(n,l,-1,index=True)

        if jss1 in out: ssw-=out[jss1]
        if jss2 in out: ssw-=out[jss2]
        if jss2 not in out:     # this subshell holds more electrons than jss1
          out[jss2]=0
          if ssw >= rlf_jssoccupancy(jss2):
            out[jss2]=rlf_jssoccupancy(jss2)
            ssw-=out[jss2]
        if jss1 not in out:
          out[jss1]=0
          if ssw > 0:
            out[jss1]=rlf_jssoccupancy(jss1)
            ssw-=out[jss1]
        if ssw != 0:
          msg='given sub-shell occupacy incompatible with given '+\
              'configuration: %s -- %s'%(inconf,ss)
          raise ErrorRAP(msg)          

    # fill in closed sub-shells
    issmin=min(out.keys())
    for iss in range(1,issmin): out[iss]=rlf_jssoccupancy(iss)

    # complete electron count
    if ne is not None:
      netot=sum(out.values())
      if netot > ne:
        msg='configuration has too many electrons, %s'%inconf
        raise ErrorRAP(msg)
      elif netot < ne:
        neleft=ne-netot
        ssmax=max(out.keys())
        for iss in range(1,ssmax):
          if neleft <= 0: break
          if iss in out: continue
          w=rlf_jssoccupancy(iss)
          if w <= neleft:
            out[iss]=w
            neleft-=w
            
        nechk=sum(out.values())
        if nechk != ne:
          msg='could not match given number of electrons'
          raise ErrorRAP(msg)

    # remove empty final subshells
    while 1:
      sss=sorted(out.keys())
      if out[sss[-1]] > 0: break
      del out[sss[-1]]

    return out

# ------------------------------------------------------------------------------

  def copy(self):
    """ Return a copy. """
    return copy.deepcopy(self)

# ------------------------------------------------------------------------------

  def label(self,showclosed=False):
    if self.num_electrons() == 0: return ''
    out=[]
    if showclosed:
      for ss in self.subshells(): out.append(ss+str(self.occupancy(ss)))
    else:
      imax=max(self.v_conf.keys())
      icore=1
      while icore in self.v_conf and \
            self.v_conf[icore] == rlf_jssoccupancy(icore):
        icore+=1
      if icore > imax: 
        icore=imax
      elif icore not in self.v_conf:
        out.append(rlf_jsubshell(icore)+'0')
      for i in self.subshells(indices=True):
        w=self.occupancy(i)
        if i < icore: continue
        out.append(rlf_jsubshell(i)+str(w))
    out=' '.join(out)
    return out

# ------------------------------------------------------------------------------

  def num_electrons(self):
    """ Return total number of electrons. """
    return sum(self.v_conf.values())

# ------------------------------------------------------------------------------

  def subshells(self,indices=False):
    """ Return sorted list of occupied sub-shells. """
    ss=self.v_conf.keys()
    ss.sort()
    if indices: return ss
    return map(rlf_jsubshell,ss)

# ------------------------------------------------------------------------------

  def occupancy(self,ss):
    """ 
    Return occupancy of sub-shell, ss (which can be index or string form. 
    """
    sss=ss
    if not isinstance(sss,int): sss=rlf_jssindex(sss)
    if sss not in self.v_conf: return 0
    return self.v_conf[sss]

# ------------------------------------------------------------------------------

  def parity(self):
    """ Return parity of configuration. """
    p=1
    for i in self.subshells():
      (n,l,s)=rlf_jssdecode(i)
      if l%2 == 0: continue
      p*=(-1)**self.occupancy(i)
    return (1-p)/2

# ------------------------------------------------------------------------------

  def weight(self):
    """ Return statistical weight of configuration. """
    wgt=1
    for ss in self.subshells():
      w=self.occupancy(ss)
      wmax=rlf_jssoccupancy(ss)
      wgt*=rlf_binom(wmax,w)
    return int(wgt)

# ------------------------------------------------------------------------------

  def valence_subshell(self):
    """ Return valence sub-shell. """
    tss=max(self.subshells(indices=True))
    return rlf_jsubshell(tss)

# ------------------------------------------------------------------------------

  def pop_electron(self,ss=None):
    """ 
    Return a copy of the configuration where an electron from the given 
    sub-shell has been removed (last sub-shell by default).
    """
    out=copy.deepcopy(self)
    tss=ss
    if tss is None: 
      tss=max(out.subshells(indices=True))
    else:
      if not isinstance(tss,int): tss=rlf_jssindex(tss)
    if tss not in out.subshells(indices=True): 
      msg='specified sub-shell not present: %s'%ss
      raise ErrorRAP(msg)
    out.v_conf[tss]-=1
    if out.v_conf[tss] == 0: del out.v_conf[tss]
    return out

# ------------------------------------------------------------------------------

  def add_electron(self,ss=None):
    """ 
    Return a copy of the configuration where an electron has been added to the
    given sub-shell (valence sub-shell by default).
    """
    out=copy.deepcopy(self)
    tss=ss
    if tss is None: 
      tss=max(out.subshells(indices=True))
    else:
      if not isinstance(tss,int): tss=rlf_jssindex(tss)
    if tss not in out.subshells(indices=True):
      out.v_conf[tss]=1
    else:
      if out.v_conf[tss] < rlf_jssoccupancy(tss):
        out.v_conf[tss]+=1
      else:
        if ss is None:
          out.v_conf[tss+1]=1
        else:
          msg='specified sub-shell is full: %s'%ss
          raise ErrorRAP(msg)
    return out

# ------------------------------------------------------------------------------

  def difference(self,conf):
    """ 
    Return dictionary of sub-shells and difference of occupation numbers.
    The sign of the difference is according to self - conf.
    """
    sss=self.subshells(indices=True)
    osss=conf.subshells(indices=True)
    out={}
    for ss in sss:
      w=self.occupancy(ss)
      if ss not in osss:
        out[rlf_jsubshell(ss)]=w
      else:
        dw=w-conf.occupancy(ss)
        if dw != 0: out[rlf_jsubshell(ss)]=dw
    for ss in osss:
      if ss in sss: continue
      out[rlf_jsubshell(ss)]=-conf.occupancy(ss)
    return out

# ------------------------------------------------------------------------------

  def to_config(self):
    """ 
    Convert configuration to rap_config form.
    """
    tmp={}
    for jss in self:
      ss=jss[:-1]
      if ss not in tmp: tmp[ss]=0
      tmp[ss]+=self[jss]
    keys=tmp.keys()
    keys.sort(rlf_sssorter)
    conf=[ss+str(tmp[ss]) for ss in keys]
    conf=' '.join(conf)
    return rap_config(conf)

